#! /usr/bin/env python3
# -*-coding: utf-8-*-

"""
Created on Apr 20

@author: Moonkie

@attention: 
"""

import mxnet as mx
import numpy as np
import cv2
import os


CURDIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.split(os.path.dirname(os.path.abspath(__file__)))[0]
DBDIR = os.path.join(CURDIR,'data')
ERROR_LOG_FILE = os.path.join(DBDIR,'error.log')
KVSTORETYPE = 'local'
LST = os.path.join(CURDIR,'train.lst')



if not os.path.exists(DBDIR):
    os.makedirs(DBDIR)

class NetGenerator(object):
    @staticmethod
    def convolutionPoolingFactory(data,name,num_filter,cov_kernel,cov_stride,pool_kernel,act_type='relu',pool_type='avg',pool_stride=(1,1),cov_pad=(0,0),pool_pad=(0,0),pool_global=False):
        cov_name = 'cov_ %s' % str(name)
        act_name = 'act_ %s' % str(name)
        pool_name = 'pool_ %s' % str(name)
        
        cov = mx.sym.Convolution(data=data,name=cov_name,kernel=cov_kernel,stride=cov_stride,pad=cov_pad,num_filter=num_filter)
        act = mx.sym.Activation(data=cov,name=cov_name,act_type=act_type)
        pool = mx.sym.Pooling(data=act,name=pool_name,kernel=pool_kernel,stride=pool_stride,pad=pool_pad,global_pool=pool_global,pool_type=pool_type)
        return pool

    @staticmethod
    def convolutionFactory(data,name,cov_kernel,cov_stride,act_type='relu',cov_pad=(0,0),cov_cudnn=False):
        cov_name = 'cov_%s' % str(name) 
        act_name = 'cov_%s' % str(name) 
        cov = mx.sym.Convolution(data=data,name=cov_name,kernel=cov_kernel,stride=cov_kernel,pad=cov_pad,cudnn_off=cov_cudnn)
        act = mx.sym.Activation(data=cov,name=act_name,act_type=act_type)
        return act

    @staticmethod
    def fullyConnectedActivationFactory(data,name,num_hidden,act_type='relu'):
        fc_name = 'fc_%s' % str(name)
        act_name = 'act_%s' % str(name)
        fc = mx.sym.FullyConnected(data=data,name=fc_name,num_hidden=num_hidden)
        act = mx.sym.Activation(data=fc,name=act_name,act_type=act_type)
        return act

    @staticmethod
    def fcActivationFcFactory(data,name,num_hidden1,num_hidden2,act_type='relu'):
        fc1_name = 'fc_%s_1' % name
        fc2_name = 'fc_%s_2' % name
        act_name = 'act_%s' % name
        fc1 = mx.sym.FullyConnected(data=data,name=fc1_name,num_hidden=num_hidden1)
        act = mx.sym.Activation(data=fc1,name=act_name,act_type=act_type)
        fc2 = mx.sym.FullyConnected(data=act,name=fc2_name,num_hidden=num_hidden2)
        return fc2


class Bpnn(object):
    '''Bpnn neural network based on mxnet 

    '''
    learning_rate = 0
    wd = 0
    momentum = 0
    
    def __init__(self,num_classes,num_feature,learn=.1,wd=0.00001,momentum=0.9,num_batch=10,ctx=mx.cpu()):
        data = mx.sym.Variable('data')
        cov1 = NetGenerator.convolutionPoolingFactory(data,name='gen1',num_filter=28, cov_kernel=(3,3),cov_stride=(1,1),pool_kernel=(3,3),pool_type='max')
        dp1 = mx.sym.Dropout(cov1,name='dp1')
        cov2 = NetGenerator.convolutionPoolingFactory(dp1,name='gen2',num_filter=32,cov_kernel=(3,3),cov_stride=(1,1),pool_kernel=(3,3),pool_type='max')
        fc1 = NetGenerator.fullyConnectedActivationFactory(cov2,name='gen3',num_hidden=128)
        dp2 = mx.sym.Dropout(fc1,name='dp2')
        fc2 = mx.sym.FullyConnected(dp2,num_hidden=num_classes,name='fc2')
        out = mx.sym.SoftmaxOutput(fc2,name='out')

        self.batch = num_batch
        self.classes = num_classes
        self.features = num_feature

        self.learning_rate = learn
        self.wd = wd
        self.momentum = momentum
        self.network = out

        self.model = mx.mod.Module(out,data_names=['data'],label_names=['out_label'],context=ctx)
        
    
    def forward(self,batch_input,is_train=False,ctx=mx.cpu()):
        self.exe.forward(is_train,data=batch_input)

    def propagation(self,out_label):
        self.args['out_label'] = out_label
        self.exe.backward()

        for weight,grad in zip(self.exe.arg_arrays,self.exe.grad_arrays):
            weight[:] -= self.learn * (grad/self.batch)
        # err = (mx.nd.argmax_channel(self.exe.outputs[0]).asnumpy() == out_label).sum()
        err = .5*(out_label - self.exe.outputs[0].asnumpy())**2
        return err

    def train(self,samples,labels,epoch=1000,is_train=True):
        parts = int(len(samples)/self.batch)
        for e in range(epoch):
            error = []
            for i in range(parts+1):
                batch_input = samples[i:i+self.batch]
                batch_label = labels[i:i+self.batch]
                self.forward(batch_input=batch_input,is_train=is_train)
                err = self.propagation(out_label=batch_label)
                error.extend(err)
            error = np.array(error)
            print('[^_^] epoch %d error is: %s' % (e,str(np.mean(error))))
        print('[^_^] Finish trainning!')
        print(self.args['out_label'],'\n','\n',self.exe.outputs[0].asnumpy())

    def visualize(self):
        mx.viz.plot_network(self.network).view(Bpnn.__name__)
    
    @staticmethod
    def extract(lstfile,num_label,encode='utf8'):
        pass
    
    @staticmethod
    def binaryzation(img):
        kernel = np.ones((2,2),np.uint8)
        if len(img.shape) == 3:
            img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
        img = cv2.blur(img,(1,1))
        img = cv2.morphologyEx(img,cv2.MORPH_OPEN,kernel,iterations=2)
        img = cv2.threshold(img,70,255,cv2.THRESH_BINARY+cv2.THRESH_TRUNC)[1]
        thresh = np.mean(img)
        img[img>thresh] += 100
        img = cv2.equalizeHist(img)
        return img

    @staticmethod
    def split_label_data(data):
        assert len(data) != 0,"A empty array is given"
        assert  len(data[0]) == 2,"Need 2d array."
        labels = data[:,0]
        dat = data[:,1]
        print('[^_^] Split labels and data sunccessfully~')
        return np.array(list(labels),dtype=np.float32),np.array(list(dat))

    @staticmethod
    def getInitialParams(data,num_classes=7,num_batch=10):
        assert len(data.shape) == 2,"Need shape 2-d array."
        _,num_feature =  data.shape
        return num_classes,num_feature,num_batch

        
if __name__ == '__main__':
    lst = os.path.join(CURDIR,'train.lst')
    total_data = Bpnn.extract(lst,7)    
    total_data = Bpnn.trunck(total_data)
    # total_data = Bpnn.classify_data_with_label(total_data)
    labels,data = Bpnn.split_label_data(total_data)

    num_classes,num_feature,num_batch = Bpnn.getInitialParams(data)
    bp = Bpnn(num_classes=7,num_feature=num_feature,num_batch=num_batch,ctx=mx.gpu())
    # bp.visualize()
    # bp.train(data,labels)
    pass


